import numpy as np
import pandas as pd
from scipy.stats import pearsonr
import matplotlib.pyplot as plt
from datetime import datetime, timedelta
import os
import multiprocessing
from tqdm import tqdm
import warnings

warnings.filterwarnings('ignore')


# 1. 数据生成和保存
def generate_data(stock_num=50, days=1000, save_path='factor_data.xlsx'):
    """
    生成随机股票数据
    :param stock_num: 股票数量
    :param days: 交易日数量
    :param save_path: 保存路径
    :return: 包含所有数据的DataFrame
    """
    np.random.seed(42)

    # 生成日期
    dates = pd.date_range(end=datetime.today(), periods=days)

    # 生成股票代码
    stocks = [f'STK_{i:03d}' for i in range(1, stock_num + 1)]

    # 生成基础数据
    data = []
    for date in dates:
        for stock in stocks:
            # 生成随机价格和交易量
            close = np.random.uniform(10, 100)
            open_price = close * np.random.uniform(0.95, 1.05)
            high = max(open_price, close) * np.random.uniform(1, 1.1)
            low = min(open_price, close) * np.random.uniform(0.9, 1)
            volume = np.random.randint(10000, 1000000)

            # 生成一些基本面数据
            pe = np.random.uniform(5, 30)
            pb = np.random.uniform(0.5, 5)
            market_cap = np.random.uniform(1e8, 1e10)

            data.append([date, stock, open_price, high, low, close, volume, pe, pb, market_cap])

    # 创建DataFrame
    columns = ['date', 'stock', 'open', 'high', 'low', 'close', 'volume', 'pe', 'pb', 'market_cap']
    df = pd.DataFrame(data, columns=columns)

    # 计算收益率 (未来1天)
    df['return'] = df.groupby('stock')['close'].pct_change().shift(-1)

    # 保存到Excel
    df.to_excel(save_path, index=False)
    print(f"数据已保存到 {save_path}")

    return df


# 2. 因子计算和评价
class FactorCalculator:
    def __init__(self, data):
        self.data = data.copy()

    def calculate_factors(self):
        """计算10个因子"""
        # 按股票分组计算
        grouped = self.data.groupby('stock')

        # 1. 动量因子 (过去5天收益率)
        self.data['momentum_5'] = grouped['close'].pct_change(5)

        # 2. 反转因子 (过去20天收益率)
        self.data['reversal_20'] = -grouped['close'].pct_change(20)

        # 3. 波动率因子 (过去20天收益率标准差)
        self.data['volatility_20'] = grouped['close'].pct_change().rolling(20).std()

        # 4. 成交量变化率 (过去5天成交量变化)
        self.data['volume_change_5'] = grouped['volume'].pct_change(5)

        # 5. PE比率
        self.data['pe_ratio'] = self.data['pe']

        # 6. PB比率
        self.data['pb_ratio'] = self.data['pb']

        # 7. 市值因子 (取对数)
        self.data['size_factor'] = np.log(self.data['market_cap'])

        # 8. 流动性因子 (成交金额/市值)
        self.data['liquidity'] = (self.data['volume'] * self.data['close']) / self.data['market_cap']

        # 9. 高低波动 (过去5天(high-low)/close)
        self.data['hl_volatility'] = (self.data['high'] - self.data['low']) / self.data['close']

        # 10. RSI近似 (过去5天上涨日比例)
        ret = grouped['close'].pct_change()
        up_days = ret.rolling(5).apply(lambda x: (x > 0).mean())
        self.data['rsi_5'] = up_days
        v = self.data
        return self.data

    @staticmethod
    def calculate_ic(factor, forward_return):
        """计算信息系数(IC)"""
        valid_idx = (~np.isnan(factor)) & (~np.isnan(forward_return))
        if sum(valid_idx) < 2:
            return np.nan
        return pearsonr(factor[valid_idx], forward_return[valid_idx])[0]

    @staticmethod
    def calculate_ir(ic_series):
        """计算信息比率(IR)"""
        return np.nanmean(ic_series) / np.nanstd(ic_series)

    def evaluate_factors(self):
        """评估因子"""
        ic_results = {}
        ic_series = {}

        factor_cols = [col for col in self.data.columns if col not in
                       ['date', 'stock', 'open', 'high', 'low', 'close', 'volume',
                        'pe', 'pb', 'market_cap', 'return']]

        for factor in factor_cols:
            ic = self.data.groupby('date').apply(
                lambda x: self.calculate_ic(x[factor], x['return']))
            ic_results[factor] = {
                'IC Mean': np.nanmean(ic),
                'IC Std': np.nanstd(ic),
                'IR': self.calculate_ir(ic),
                'IC>0 Ratio': (ic > 0).mean()
            }
            ic_series[factor] = ic

        ic_df = pd.DataFrame(ic_results).T
        return ic_df, ic_series


# 3. 投资策略接口
class FactorStrategy:
    def __init__(self, data, factor_weights=None):
        """
        初始化策略
        :param data: 包含因子和收益率的数据
        :param factor_weights: 因子权重字典，如果为None则等权
        """
        self.data = data.copy()
        self.factor_cols = [col for col in self.data.columns if col not in
                            ['date', 'stock', 'open', 'high', 'low', 'close',
                             'volume', 'pe', 'pb', 'market_cap', 'return']]

        if factor_weights is None:
            self.factor_weights = {factor: 1 / len(self.factor_cols) for factor in self.factor_cols}
        else:
            self.factor_weights = factor_weights

        # 标准化因子
        self._standardize_factors()

    def _standardize_factors(self):
        """横截面标准化因子"""
        for factor in self.factor_cols:
            self.data[factor] = self.data.groupby('date')[factor].transform(
                lambda x: (x - x.mean()) / x.std())

    def calculate_composite_score(self):
        """计算综合因子得分"""
        self.data['composite_score'] = 0
        for factor, weight in self.factor_weights.items():
            self.data['composite_score'] += self.data[factor] * weight
        return self.data

    def generate_signals(self, top_n=10):
        """生成交易信号
        :param top_n: 买入前top_n只股票
        :return: 包含信号的DataFrame
        """
        self.calculate_composite_score()

        # 按日期分组，对综合得分排序
        self.data['rank'] = self.data.groupby('date')['composite_score'].rank(ascending=False)

        # 生成信号: 1表示买入，0表示不持有
        self.data['signal'] = (self.data['rank'] <= top_n).astype(int)

        return self.data

    def get_portfolio(self, date):
        """获取特定日期的投资组合
        :param date: 日期
        :return: 该日期的投资组合(股票列表)
        """
        daily_data = self.data[self.data['date'] == date]
        return daily_data[daily_data['signal'] == 1]['stock'].tolist()


# 4. 回测框架
class Backtest:
    def __init__(self, data, strategy, initial_capital=1e6, top_n=10):
        """
        初始化回测
        :param data: 包含信号和收益率的数据
        :param strategy: 策略对象
        :param initial_capital: 初始资金
        :param top_n: 买入前top_n只股票
        """
        self.data = data.copy()
        self.strategy = strategy
        self.initial_capital = initial_capital
        self.top_n = top_n
        self.results = None

    def run(self):
        """运行回测"""
        # 生成信号
        self.data = self.strategy.generate_signals(top_n=self.top_n)

        # 按日期进行回测
        dates = sorted(self.data['date'].unique())
        portfolio_values = []
        current_capital = self.initial_capital
        positions = {}  # {stock: shares}

        for i, date in enumerate(dates):
            # 获取当天的信号
            daily_data = self.data[self.data['date'] == date]

            # 卖出当前持仓
            for stock in list(positions.keys()):
                if stock in daily_data['stock'].values:
                    stock_data = daily_data[daily_data['stock'] == stock].iloc[0]
                    current_capital += positions[stock] * stock_data['close']
                    del positions[stock]

            # 买入新信号
            buy_stocks = daily_data[daily_data['signal'] == 1]
            if len(buy_stocks) > 0:
                # 等权分配资金
                capital_per_stock = current_capital / len(buy_stocks)
                for _, row in buy_stocks.iterrows():
                    shares = capital_per_stock / row['close']
                    positions[row['stock']] = shares
                current_capital = 0

            # 计算组合价值
            portfolio_value = current_capital
            for stock, shares in positions.items():
                if stock in daily_data['stock'].values:
                    stock_close = daily_data[daily_data['stock'] == stock]['close'].values[0]
                    portfolio_value += shares * stock_close

            portfolio_values.append(portfolio_value)

        # 保存回测结果
        self.results = pd.DataFrame({
            'date': dates,
            'portfolio_value': portfolio_values
        })

        return self.results

    def evaluate_performance(self):
        """评估回测表现"""
        if self.results is None:
            raise ValueError("请先运行回测")

        results = self.results.copy()
        results['daily_return'] = results['portfolio_value'].pct_change()

        # 计算年化收益率
        total_return = (results['portfolio_value'].iloc[-1] / self.initial_capital) - 1
        days = len(results)
        annualized_return = (1 + total_return) ** (252 / days) - 1

        # 计算年化波动率
        annualized_vol = results['daily_return'].std() * np.sqrt(252)

        # 计算夏普比率
        sharpe_ratio = annualized_return / annualized_vol if annualized_vol != 0 else np.nan

        # 计算最大回撤
        results['peak'] = results['portfolio_value'].cummax()
        results['drawdown'] = (results['portfolio_value'] - results['peak']) / results['peak']
        max_drawdown = results['drawdown'].min()

        # 计算胜率
        win_rate = (results['daily_return'] > 0).mean()

        return {
            'Total Return': total_return,
            'Annualized Return': annualized_return,
            'Annualized Volatility': annualized_vol,
            'Sharpe Ratio': sharpe_ratio,
            'Max Drawdown': max_drawdown,
            'Win Rate': win_rate,
            'Top N': self.top_n
        }

    def plot_results(self):
        """绘制回测结果"""
        if self.results is None:
            raise ValueError("请先运行回测")

        plt.figure(figsize=(12, 6))
        plt.plot(self.results['date'], self.results['portfolio_value'])
        plt.title('Portfolio Value Over Time')
        plt.xlabel('Date')
        plt.ylabel('Portfolio Value')
        plt.grid(True)
        plt.show()


# 5. 并行回测和参数优化
def parallel_backtest(params, data):
    """并行执行的回测函数"""
    top_n, factor_weights = params
    strategy = FactorStrategy(data, factor_weights)
    backtest = Backtest(data, strategy, top_n=top_n)
    backtest.run()
    performance = backtest.evaluate_performance()
    return performance


def optimize_parameters(data, top_n_list=None, factor_combinations=None, n_jobs=-1):
    """
    优化策略参数
    :param data: 数据
    :param top_n_list: 要测试的top_n列表
    :param factor_combinations: 要测试的因子组合列表
    :param n_jobs: 并行工作数
    :return: 最佳参数和结果
    """
    if top_n_list is None:
        top_n_list = [5, 10, 15, 20]

    if factor_combinations is None:
        # 默认测试所有因子等权组合
        factor_cols = [col for col in data.columns if col not in
                       ['date', 'stock', 'open', 'high', 'low', 'close',
                        'volume', 'pe', 'pb', 'market_cap', 'return']]
        factor_combinations = [None]  # None表示等权

    # 生成参数组合
    param_combinations = [(top_n, weights) for top_n in top_n_list
                          for weights in factor_combinations]

    # 并行回测
    if n_jobs == -1:
        n_jobs = multiprocessing.cpu_count()

    with multiprocessing.Pool(n_jobs) as pool:
        results = list(tqdm(pool.imap(
            lambda params: parallel_backtest(params, data), param_combinations),
            total=len(param_combinations)))

    # 找到最佳参数
    best_result = max(results, key=lambda x: x['Sharpe Ratio'])
    best_idx = results.index(best_result)
    best_params = param_combinations[best_idx]

    return best_params, results


# 主函数
def main():
    # 1. 生成或加载数据
    data_file = 'factor_data.xlsx'
    if not os.path.exists(data_file):
        print("生成新数据...")
        data = generate_data(save_path=data_file)
    else:
        print("加载已有数据...")
        data = pd.read_excel(data_file)

    # 2. 计算因子并评估
    print("\n计算因子...")
    factor_calc = FactorCalculator(data)
    data_with_factors = factor_calc.calculate_factors()

    print("\n评估因子...")
    ic_df, ic_series = factor_calc.evaluate_factors()
    print("\n因子评价指标:")
    print(ic_df)

    # 3. 创建策略
    print("\n创建策略...")
    strategy = FactorStrategy(data_with_factors)

    # 4. 运行回测
    print("\n运行回测...")
    backtest = Backtest(data_with_factors, strategy, top_n=10)
    backtest.run()
    performance = backtest.evaluate_performance()

    print("\n回测表现:")
    for k, v in performance.items():
        print(f"{k}: {v:.4f}")

    backtest.plot_results()

    # 5. 参数优化
    print("\n参数优化...")
    # 测试不同的top_n值和因子组合
    top_n_list = [5, 10, 15, 20]

    # 定义几个不同的因子权重组合
    factor_cols = [col for col in data_with_factors.columns if col not in
                   ['date', 'stock', 'open', 'high', 'low', 'close',
                    'volume', 'pe', 'pb', 'market_cap', 'return']]

    # 组合1: 动量因子权重加倍
    weights1 = {factor: 1 / len(factor_cols) for factor in factor_cols}
    weights1['momentum_5'] *= 2

    # 组合2: 只使用前5个因子
    weights2 = {factor: 0 for factor in factor_cols}
    for factor in factor_cols[:5]:
        weights2[factor] = 1 / 5

    factor_combinations = [None, weights1, weights2]

    best_params, all_results = optimize_parameters(
        data_with_factors, top_n_list=top_n_list,
        factor_combinations=factor_combinations)

    print("\n最佳参数:")
    print(f"Top N: {best_params[0]}")
    print("因子权重:")
    if best_params[1] is None:
        print("等权")
    else:
        for factor, weight in best_params[1].items():
            print(f"{factor}: {weight:.4f}")

    # 找到最佳结果
    best_result = max(all_results, key=lambda x: x['Sharpe Ratio'])
    print("\n最佳回测表现:")
    for k, v in best_result.items():
        print(f"{k}: {v:.4f}")


if __name__ == "__main__":
    main()